Draft PR for observability.

Ports the core model into a Keras model.
VAE and rest of the pipeline would make sense in a separate PR.

Notebook with numerical comparisons of modules: https://colab.research.google.com/drive/1TUI2outgEXA1-dczw4SU8i6z_7oj4CGu?usp=sharing

Includes:

• Maths module
  • Timestep embedding
  • RoPE
  • Attention
  • Scaled dot product attention re-implementation in Keras (to match the PyTorch one)
• Layers module
  • MLPEmbedder
  • RMSNorm
  • QKNorm
  • SelfAttention
  • Modulation
  • DoubleStreamBlock
  • SingleStreamBlock

TODO:

• Combine components into model
• Expand on documentation
• Verify all components having numerically equivalent behavior through demo notebook